feat(qwen36): routed experts from a qpack container through bounded Metal slots - #1323

load_tokenizer indexed the BPE merge table from tokenizer.json by
splitting each entry on a space: the legacy "a b" string spelling.
tokenizers >= 0.20 (transformers 4.45+) writes each merge as a
two-element array ["a","b"] instead, and that is what the Qwen3.6
checkpoints ship. Every entry failed the string test, `continue` skipped
it, and the reader finished with ZERO merges and no message.

Nothing downstream treats an empty merge table as an error: bpe_piece
reads it as "nothing to merge", so encode_text degraded to one token per
byte-symbol. Measured on a real-model run before the fix: a 24-character
prompt became 24 tokens, and the model was fed a sequence it had never
seen in training.

Accept both spellings. A string entry splits on the first space as
before; a two-element array of strings takes its halves directly; any
other shape is skipped as it was. Both forms index into the same
smap_put(&g_merge, key, r) table with the same 0x1F-joined key, so the
rest of the tokenizer is unchanged.

tests/test_qwen36_tok_merges writes two tokenizer.json files for the
same five-piece vocabulary that differ only in the merge spelling and
asserts that both encode "in in" to the same merged ids [in, Ġin] -- one
plain merge (i+n) and one chained merge (Ġ+in) whose input the first
produces. On the pre-fix reader the pair form loads zero merges and the
encoding comes back as one id per byte-symbol, so the gate fails. The
scratch directories the test binaries write (tests/tmp_*) are added to
.gitignore.

(cherry picked from commit d783505)

The Darwin arm of the Makefile finds libomp with `brew --prefix libomp`.
On a CLT-only Mac driven non-interactively (`make` from a script, over
SSH, in CI) Homebrew is routinely installed but not on PATH, the probe
prints nothing, and the OpenMP arm never engages. Nothing warns: the
build is a valid single-threaded engine with libomp sitting right there
on disk.

Measured on an M4 Max with the Qwen3.6-35B qpack: 0.87 tok/s from the
silent single-threaded build, 4.44 tok/s (range 3.80-4.76, thermal) once
OpenMP engaged, TTFT 3.89 s to about 0.7 s, output byte-identical by
cmp. Engine code is unchanged by this commit -- every omp site is
row-partitioned and deterministic by construction -- so the difference
is entirely the missing flag.

When the brew probe comes back empty, fall back to the two standard
Homebrew prefixes (/opt/homebrew/opt/libomp on Apple silicon,
/usr/local/opt/libomp on Intel). The existing artifact check still runs
after the fallback: both omp.h and libomp.* must exist before any flag
is added, so a bare prefix directory without the library still yields
the plain build, as before.

(cherry picked from commit 44ddb39)

…ract

affine_quant.h carries one reading of MLX's packed affine words:
coli_affine_matmul_ref folds scale*q + bias into its accumulation and
never materialises the weights. That is the right shape for a routed
expert dispatched per token, and the wrong shape for the dense tensors
the qwen36 engine keeps RESIDENT -- attention, DeltaNet, shared expert,
norms, embeddings -- which a Swiftlet qpack container also stores as
affine triples in its model.safetensors and which the engine wants as
f32 rows exactly once, at load.

coli_affine_dequant_ref expands a validated view into row-major
out[output_dim, input_dim] = scale[row, col/gs] * q[row, col] +
bias[row, col/gs], with the same bit conventions as matmul_ref
(little-endian uint32 words, lowest bits = lowest logical column, Q4 or
Q8, BF16/F16/F32 scalars). It runs coli_affine_validate first, so a
malformed view comes back as a status and never as plausible output. It
lives next to matmul_ref deliberately: the two bit-level readings of the
same words cannot drift apart without this file changing.

tests/test_affine_dequant pins the contract from the bits up: a Q4 word
authored as 0x76543210 must expand to logical columns 0..7 in that order
(a most-significant-first reading produces the reversed sequence and
fails the exact-value checks); the affine map is checked with values
exactly representable in BF16, F16 and F32 so the comparisons are
equalities, not tolerances; a crafted Q8 word proves the byte order; on
a multi-group view x @ dequant(W)^T must equal matmul_ref(x, W); and
NULL buffers, truncated scales, and a group size that does not divide
input_dim are refused with statuses.

A Swiftlet qpack container carries the dense half of the model in its
model.safetensors exactly as mlx-lm wrote it: every quantized tensor is
a packed U32 `.weight` with `.scales`/`.biases` siblings (BF16 on the
production container), and the text stack sits under a `language_model.`
prefix because the checkpoint is multimodal. qwen36's loader could do
nothing with that file -- st_init exit(1)'d on the first U32 dtype
string before the rest of the header was even indexed, and load_t_n only
speaks unprefixed names. Running the engine over a container therefore
meant re-converting weights whose bytes the container already holds.

Three additive changes let the loader read the container as it is:

  st.h indexes U32/I32 as dtype 7 (4 bytes/element), the same way the
  fp8 codes were added: the header parses to the end, the tensor is
  findable with its packed shape, and the float readers keep refusing it
  by value (dtype >= 3 in st_read_f32) -- callers must go through
  st_read_raw with geometry they validated.

  dense_resolve tries the plain name first and `language_model.<name>`
  second, so converted snapshots are untouched and the container's
  spelling is found; when neither exists it returns the CANONICAL name
  so the refusal names the tensor the engine wanted.

  load_t_affine expands a U32 triple through coli_affine_dequant_ref.
  Every dimension is anchored to `want`, the config-implied element
  count the forward pass will index with -- the discipline load_t_n
  already applies to plain tensors -- so a container whose packed
  geometry disagrees with config.json is a refusal, never a plausible
  heap OOB. Bits (Q4 = 8 per word, Q8 = 4) and the group size are
  DERIVED from the shapes (logical input over packed words, logical
  input over scale groups) rather than parsed from config quantization
  overrides: the file's own byte layout is what the expansion must agree
  with, and coli_affine_validate re-checks the derived geometry against
  every buffer length. Scales in BF16/F16/F32 are accepted; any other
  dtype, mismatched biases, an orphan `.weight` without siblings, a
  `want` the packed shape cannot tile, or a per-word count that is
  neither 4 nor 8 exits with a message naming the tensor.

load_t_n dispatches to load_t_affine on dtype 7 and is otherwise
unchanged; the optional-tensor probes (q/k norms, router bias, shared
expert gate) go through dense_has so the prefixed spelling is found too.
The router correction bias now loads through load_t_n as well, which
gives it the size check the other dense tensors already had.

tests/test_qwen36_dense_affine builds a safetensors file by hand (no
numpy, no torch): both name spellings load and the plain one wins when
both exist; st_init indexes the U32 tensors with the right rank and
shape while st_read_f32 still refuses them; a Q4 and a Q8 triple with
authored bit patterns expand to exactly scale*q + bias; and (POSIX fork
gates) an untileable want, an orphan U32 weight, and a wrong-size plain
tensor each exit non-zero instead of returning rows. With argv the same
binary loads one named tensor through load_t_n and compares it against
an expected.f32 file, which the snap-view tool test uses as its
cross-implementation round trip.

…d form

rmsnorm_row applies its weight as (1 + w). That is the HF Qwen3.6
convention -- the norm weights are stored zero-centered -- and every
snapshot convert_qwen36.py produces keeps it, so the forward pass was
written for it and has only ever seen it.

mlx-lm materialises the +1 into the weights at conversion, so an
MLX-derived container stores FULL gamma. Measured on the production
Qwen3.6-35B qpack: input_layernorm mean 1.03, q_norm mean 1.33, where
the zero-centered forms centre on 0. Fed through (1 + w) that doubles
every normalised activation and the model degenerates to noise, with
every tensor loaded and every shape check passing.

The shift is undone at LOAD, so the forward pass keeps exactly one
convention. Cfg gains zero_centered_norms (default 1, the HF dialect;
qwen36_meta.json may set it false, which is what the snap-view tool
emits for a container). load_norm_n loads through load_t_n and, when the
flag is 0, subtracts 1 from every element. It is used for precisely the
tensors rmsnorm_row touches: input/post layernorms, q/k norms, and the
final norm in both the main loader and the edge engine. The MoE router
weight, which shared the LD macro with the layernorms, is loaded
separately and never shifted. The DeltaNet gated norm is full gamma in
BOTH dialects (its forward multiplies plain w) and does not go through
this loader.

tests/test_qwen36_dense_affine gains two gates: load_norm_n leaves the
values untouched under the default dialect and returns exactly value-1
under the MLX one, and load_meta parses a zero_centered_norms:false
qwen36_meta.json into the config.

The engine's loading contract is the one convert_qwen36.py established:
a snapshot directory with a FLAT config.json, a qwen36_meta.json whose
head dims were derived from the actual weight shapes, the *.safetensors
holding the dense weights, and tokenizer.json. A Swiftlet qpack
container already carries every byte the engine needs --
model.safetensors (the dense half, which the loader now expands in
memory), tokenizer.json, and a nested multimodal config.json -- but not
the two engine-side metadata files, and the container itself must never
be written: it is a production artifact shared with Swiftlet and
hash-pinned by hashes.json.

So the bridge is a VIEW. tools/make_qwen36_qpack_snap.py writes a fresh
directory holding a generated flat config.json and qwen36_meta.json next
to symlinks into the container. No tensor bytes are copied or converted
on disk; st_init follows the link and the affine dense loader does the
expansion at load. Stdlib only: the safetensors header is 8 bytes of
length plus JSON, so there is no torch, numpy or safetensors dependency.
An --out inside the container is refused.

Head-dim derivation mirrors convert_qwen36.py ("derived from the actual
weight shapes, authoritative") with one extra step: quantized tensors
declare PACKED shapes (uint32 words along the input axis), so logical
dims are unpacked through the config's MLX quantization spec -- default
bits/group plus the per-module overrides, matching what Swiftlet's
Checkpoint.quantSpec applies. The meta declares zero_centered_norms:
false, because an mlx-lm container stores full-gamma norm weights and
the engine unshifts them at load on that flag. The argparse description
falls back to a literal under python -OO, where __doc__ is None.

  python tools/make_qwen36_qpack_snap.py \
      --container ~/models/qwen3.6-35b.qpack --out ~/build/qwen36-35b-snap
  SNAP=<out> QWEN36_QPACK=<container> ./qwen36 16 4 prompt.txt

tests/test_make_qwen36_qpack_snap.py builds a synthetic container
(nested config with per-module overrides, a hand-written safetensors
with packed U32 triples) and checks the flattened config, the unpacked
head dims (a derivation that forgets to unpack reports o_in 8x too small
and fails), layer_types / n_active / DeltaNet dims / rope_theta
surviving the trip, the links, the untouched container, the refusal, and
the dialect flag. The round trip is the cross-implementation evidence:
the fixture carries a seeded-random Q4 tensor whose dequantized rows are
computed in pure Python (bf16 decoded exactly, scale*q + bias per group)
and handed to tests/test_qwen36_dense_affine's compare mode, which loads
the same bytes through the engine's own st_init + load_t_n path via the
view's symlink. It skips loudly if that gate binary is not built.

coli_metal_matmul_affine keys its persistent ColiMetalTensor on the
projection's host pointer: the first call wraps or copies the bytes into
an MTLBuffer and every later call with the same pointer reuses that
buffer. That contract is right for weights that never move and wrong for
memory that is REFILLED -- a bounded pool of expert slots that reads a
different expert into the same page-aligned blob on every miss. The
cache would keep serving the GPU-side snapshot of whichever expert
filled the slot first, and nothing would say so.

A slot buffer is the shape that memory needs. coli_metal_slot_register
wraps one page-aligned host slot in an MTLBuffer exactly ONCE,
zero-copy, and REFUSES anything that would not wrap in place: base must
be aligned to 16384 (the Apple page) and len a multiple of it, because
the copying fallback used elsewhere would silently detach the GPU from
later refills. qpack expert strides already satisfy this
(COLI_QPACK_PAGE_ALIGNMENT). Nothing is keyed on the refillable pointer
and no bytes are copied at registration, so a refill cannot leave a
stale snapshot behind. Slot buffers count in coli_metal_stats like
resident handles.

coli_metal_matmul_affine_slot runs the same Q4/Q8 affine kernels as
coli_metal_matmul_affine, addressing the weight, scale, and bias
sections by BYTE OFFSET inside the registered buffer (setBuffer:offset:)
instead of through a cached handle. The checked view must describe the
slot's CURRENT fill: every section must lie inside the registered range,
the weights offset must be 4-byte aligned (the kernel reads uint32
words), and the view's host pointers must equal base+offset. A
descriptor from a previous fill or the wrong slot is refused (returns 0,
the CPU-fallback contract) rather than dispatched. Dispatch is
synchronous like the rest of this API: when it returns, the GPU is done
reading the slot and the caller may refill it. The encode path is
factored into affine_dispatch_run, shared with the handle-based entry,
so the two dispatch the same kernels with the same parameters.

metal-test gains the slot-buffer arm: offset addressing inside one
whole-slot buffer matches the CPU reference for Q4 and Q8; refilling the
same registered memory with a different expert is seen by the next
dispatch (the write-through the handle cache cannot give reused memory);
and the refusals hold -- misaligned registration, sections outside the
registered range, and a view that does not describe the slot's current
fill are refused, never dispatched.

…lots

With the dense half loading from a container's model.safetensors, the
routed experts are what still kept qwen36 on a converted snapshot. This
gives the engine a second expert source: QWEN36_QPACK=<dir> opens a
Swiftlet qpack v1 container and moe() computes every routed expert from
it through the MLX affine contract -- coli_metal_matmul_affine_slot on
Apple builds, coli_affine_matmul_ref everywhere else. Router, SwiGLU
combine, shared expert, attention, and DeltaNet stay on the CPU; this is
NOT a whole-GPU forward pass. The CUDA expert tier is not initialised
when the container owns the experts (it would warmstart-load snapshot
experts moe() no longer computes), and a failed expert is fatal: there
is no other weight source to fall back to, and a silently skipped expert
would just be a wrong answer.

qwen36_qpack.[ch] is the store. qq_open validates the container against
the loaded config -- layerCount/expertCount, and expert (0,0)'s
gate/up/down projections against [inter,hidden]/[inter,hidden]/
[hidden,inter] -- and refuses a mismatch with a message; the engine
treats that refusal as fatal rather than falling back to the snapshot.
Experts live in a BOUNDED pool of refillable slots. Each slot owns one
page-aligned expert_stride blob (the qpack layout guarantees
expert_stride % 16384 == 0) registered with the Metal backend exactly
once as a whole-slot buffer; a miss refills the least-recently-used slot
in place with one pread and bumps the slot's generation. The forward is
split into qq_expert_acquire, which returns a QqExpertRef {slot,
generation}, and qq_ref_forward, which re-checks the generation at
dispatch and refuses a ref captured against an earlier fill without
touching the accumulator -- so a caller that held a ref across an
eviction gets a refusal, never another expert's bytes. Metal dispatch
addresses each projection's weight/scale/bias section by byte offset
inside the slot's registration, so nothing is keyed on the refillable
host pointer. The pool defaults to 96 slots, always clamped to the
container's expert total (tiny fixtures stay fully resident, real
containers are bounded); QWEN36_QPACK_SLOTS=<n> overrides it.
qq_slot_stats exposes slots/evictions/fills and qq_counts the per-path
projection counts, so a test can PROVE eviction happened and which path
ran instead of trusting a silent fallback.

The affine format boundary holds: these blobs are unsigned uint32-packed
MLX Q4/Q8 with scale AND bias, never Colibri's signed-int4 QT fmt=4, and
they never enter the QT fmt namespace. The store is single-threaded by
contract: only the engine's MoE loop touches it.

On Apple builds the engine calls coli_metal_init() before qq_open. The
store dispatches through coli_metal_matmul_affine_slot, which is gated
on pipelines only coli_metal_init compiles; the parity gates call it in
their own main, so an engine that did not would run every real-model
routed projection on the CPU reference while looking identical apart
from the counters (measured on the first Qwen3.6-35B smoke: metal=0
cpu=37440). Metal being unavailable is reported, not refused -- the CPU
reference is the documented oracle arm -- but it is never a silent
default.

Build wiring: METAL=1 engine builds compile qwen36_qpack.c with
-DQWEN36_QPACK; every other build keeps the inline stubs in the header,
which return 0 and a message ("engine built without qpack affine
support") so QWEN36_QPACK=<dir> on such a build is a loud refusal. This
is the same arrangement as the CUDA tier in qwen36_tier.h.

tests/test_qwen36_qpack is the CPU gate, on a container synthesized the
way tests/test_qpack.c does it (2 layers x 2 experts, Q4, group 8), no
GPU in the room: qq_open refuses each mismatched dimension and accepts
the matching container; qq_expert_forward equals an independent
composition of the affine reference (gate/up/down + SwiGLU + weighted
accumulate) and refuses out-of-range layer/expert ids with a real out
buffer, so the range check itself is what is tested; moe() with the
store active produces exactly the router-weighted sum of reference
expert outputs; a pool bounded to ONE slot serving four experts
evicts/refills on every forward (evictions=7 fills=8 asserted) with
parity to the reference; and a ref held across an evict/refill is
refused at dispatch with the accumulator untouched, while a re-acquired
ref carries a new generation and dispatches.

The CPU gate proves the store's glue against the affine reference with
no GPU in the room. This is the hardware twin: `make qwen36-metal-test`
runs the qwen36 routed experts through the Metal affine kernels against
the CPU reference on a REAL Swiftlet-generated container -- set
QWEN36_QPACK_FIXTURE=<dir> to a swiftlet-repack output of
fixtures/tiny-model-q4. Without the fixture the binary SKIPS with a
distinct marker; it never pretends to have proven parity. The logit gate
is pinned to the tiny-model-q4 geometry (the DeltaNet/attention dims are
not in the container), and any other geometry fails loudly.

The slot pool is forced SMALLER than one layer's expert population
(EVICT_SLOTS = 5 against 8 experts per layer, 64 (layer,expert) pairs
across the stack), so none of the gates can pass without live
eviction/refill of the whole-slot Metal buffers during the runs. Three
gates, the first two running the SAME engine code with the dispatch
flipped by qq_force_cpu():

  1. one-layer parity: moe() on one MoE block, routed experts through
     the Metal affine path vs the CPU affine reference, same router,
     same shared expert.
  2. logit parity: an end-to-end step() (prefill S=5 plus one decode
     token) through the full tiny hybrid stack -- DeltaNet + Gated
     Attention + MoE on every layer -- with logits compared between the
     two dispatch modes. Attention and DeltaNet run on the CPU in BOTH
     modes; only the routed experts move, which is exactly the claim.
  3. stale-ref refusal on hardware: acquire an expert, cycle the pool
     until its slot is refilled with a different expert, then present
     the old-generation ref. Dispatch must refuse it with ZERO Metal or
     CPU projections and the accumulator untouched, and a fresh acquire
     must dispatch on Metal again.

Dispatch is proven, not assumed. qq_counts must show every routed
projection of the Metal run on the GPU and none on the CPU fallback,
qq_slot_stats must show evictions actually happened during the parity
runs, and coli_metal_stats must show the resident slot buffers bounded
by the pool size (EVICT_SLOTS whole-slot buffers, where a resident
per-projection cache would hold three handles per expert).

Measured on an M4 Max against the tiny-q4 fixture: one-layer moe nerr
4.86e-07 (18/18 projections on Metal, 0 CPU fallback), logits prefill
nerr 3.58e-07, logits decode nerr 5.85e-07, routed dispatch 288/288 on
Metal, evictions=136 fills=141 over the parity runs, stale ref refused
(refused=1 accumulator-untouched=1 zero-dispatch=1), resident footprint
5 slot buffers / 81,920 bytes.

The parity gates assert which path computed the routed projections by
reading qq_counts: coli_metal_matmul_affine_slot returns 0 for a CPU
fallback by contract, and the store's counters are where that becomes
visible. The engine had no such line. A METAL=1 run whose experts all
fell back to the CPU affine reference produced the same text, the same
exit code, and the same log as one on the GPU; the only tell was the
wall clock. That is exactly how the missing coli_metal_init went
unnoticed until the counters were printed (metal=0 cpu=37440 on the
first Qwen3.6-35B smoke).

When a store is active, print its evidence next to the other exit
statistics: projections by path (metal / cpu) and how hard the bounded
slot pool worked (slots / fills / evictions). A run whose experts
silently fell back shows cpu != 0 here instead of hiding it.

Measured on an M4 Max with the production Qwen3.6-35B qpack (16 greedy
tokens after "The capital of France is", the synchronous per-expert
path): [qpack] projections: metal=19200 cpu=0 | slots=96 fills=5914
evictions=5818.